import numpy as np

def hunter_decision(obs):
    """
    猎手决策逻辑：始终追踪距离自己最近的逃脱者
    
    Args:
        obs: 包含猎手自身信息和其他智能体信息的观察数据
            - agent_obs: 猎手自身的观察数据，包含position和angle
            - other_agent_obs: 其他智能体的观察数据字典
    
    Returns:
        numpy数组: [加速度, 角速度]，表示猎手的动作
    """
    # 初始化动作，加速度为0，角速度为0°
    action = [3, 4]
    
    # 从观察数据中提取猎手位置和角度
    hunter_position = obs['agent_obs']['position']
    hunter_angle = obs['agent_obs']['angle'][0] * np.pi * 2  # 当前朝向角度（弧度）
    
    # 寻找最近的逃脱者
    min_distance = float('inf')
    nearest_escaper_position = None
    
    # 遍历所有其他智能体，找到距离最近的
    for agent_id, agent_info in obs['other_agent_obs'].items():
        # 确保智能体位置有效
        if agent_info['position'].any():
            # 计算欧几里得距离
            distance = np.linalg.norm(hunter_position - agent_info['position'])
            
            # 更新最近目标
            if distance < min_distance:
                min_distance = distance
                nearest_escaper_position = agent_info['position']
    
    # 如果找到最近的逃脱者，生成追踪动作
    if nearest_escaper_position is not None:
        # 计算指向最近逃脱者的方向向量
        direction_vector = nearest_escaper_position - hunter_position
        
        # 计算目标角度（弧度）
        target_angle = np.arctan2(direction_vector[1], direction_vector[0])
        
        # 计算角度差，并将其限制在[-π, π]范围内
        angle_diff = target_angle - hunter_angle
        angle_diff = (angle_diff + np.pi) % (2 * np.pi) - np.pi
        
        # 根据角度差选择合适的角速度动作
        if angle_diff > 0:
            # 需要向右转向
            if angle_diff < np.pi / 8:
                action[1] = 4  # 0°（不需要转向）
            elif angle_diff < 3 * np.pi / 8:
                action[1] = 5  # 22.5°
            elif angle_diff < 5 * np.pi / 8:
                action[1] = 6  # 45°
            elif angle_diff < 7 * np.pi / 8:
                action[1] = 7  # 67.5°
            else:
                action[1] = 8  # 90°（最大右转角度）
        else:
            # 需要向左转向
            if angle_diff > -np.pi / 8:
                action[1] = 4  # 0°（不需要转向）
            elif angle_diff > -3 * np.pi / 8:
                action[1] = 3  # -22.5°
            elif angle_diff > -5 * np.pi / 8:
                action[1] = 2  # -45°
            elif angle_diff > -7 * np.pi / 8:
                action[1] = 1  # -67.5°
            else:
                action[1] = 0  # -90°（最大左转角度）
        
        # 始终选择最大加速度，全力追击
        action[0] = 4
    
    # 如果没有找到逃脱者，保持当前状态
    else:
        # 保持零加速度和零角速度
        action = [3, 4]
    
    return np.array(action, dtype=np.int64)

def escaper_decision(obs):
    """
    逃脱者决策逻辑：沿猎手朝向自己的连线的垂直方向进行逃脱
    
    Args:
        obs: 包含逃脱者自身信息和其他智能体信息的观察数据
            - agent_obs: 逃脱者自身的观察数据，包含position和angle
            - other_agent_obs: 其他智能体的观察数据字典
    
    Returns:
        numpy数组: [加速度, 角速度]，表示逃脱者的动作
    """
    # 初始化动作，加速度为0，角速度为0°
    action = [3, 4]
    
    # 从观察数据中提取逃脱者位置和角度
    escaper_position = obs['agent_obs']['position']
    escaper_angle = obs['agent_obs']['angle'][0] * np.pi * 2  # 当前朝向角度（弧度）
    
    # 寻找最近的猎手（威胁）
    min_distance = float('inf')
    nearest_hunter_position = None
    
    # 遍历所有其他智能体，找到距离最近的
    for agent_id, agent_info in obs['other_agent_obs'].items():
        # 确保智能体位置有效
        if agent_info['position'].any():
            # 计算欧几里得距离
            distance = np.linalg.norm(escaper_position - agent_info['position'])
            
            # 更新最近威胁
            if distance < min_distance:
                min_distance = distance
                nearest_hunter_position = agent_info['position']
    
    # 如果找到最近的猎手，生成逃脱动作
    if nearest_hunter_position is not None:
        # 计算猎手朝向逃脱者的方向向量
        threat_direction = escaper_position - nearest_hunter_position
        
        # 规范化方向向量
        direction_norm = np.linalg.norm(threat_direction)
        if direction_norm > 0:  # 避免除以零
            unit_threat_direction = threat_direction / direction_norm
        else:
            unit_threat_direction = np.array([1.0, 0.0])  # 默认方向
        
        # 计算两个可能的垂直方向（法线方向）
        # 垂直方向1：逆时针旋转90度 (dx, dy) -> (-dy, dx)
        vertical_dir1 = np.array([-unit_threat_direction[1], unit_threat_direction[0]])
        # 垂直方向2：顺时针旋转90度 (dx, dy) -> (dy, -dx)
        vertical_dir2 = np.array([unit_threat_direction[1], -unit_threat_direction[0]])
        
        # 计算逃脱者当前朝向与两个垂直方向之间的角度
        current_direction = np.array([np.cos(escaper_angle), np.sin(escaper_angle)])
        
        # 计算逃脱者当前朝向与每个垂直方向的夹角
        angle_with_dir1 = np.arccos(np.clip(np.dot(current_direction, vertical_dir1), -1, 1))
        angle_with_dir2 = np.arccos(np.clip(np.dot(current_direction, vertical_dir2), -1, 1))
        
        # 选择当前朝向与垂直方向夹角较小的方向，这样转向更快
        if angle_with_dir1 <= angle_with_dir2:
            escape_direction = vertical_dir1
            target_angle = np.arctan2(escape_direction[1], escape_direction[0])
        else:
            escape_direction = vertical_dir2
            target_angle = np.arctan2(escape_direction[1], escape_direction[0])
        
        # 计算需要转向的角度差，并将其限制在[-π, π]范围内
        angle_diff = target_angle - escaper_angle
        angle_diff = (angle_diff + np.pi) % (2 * np.pi) - np.pi
        
        # 根据角度差选择合适的角速度动作
        if angle_diff > 0:
            # 需要向右转向
            if angle_diff < np.pi / 8:
                action[1] = 4  # 0°（不需要转向）
            elif angle_diff < 3 * np.pi / 8:
                action[1] = 5  # 22.5°
            elif angle_diff < 5 * np.pi / 8:
                action[1] = 6  # 45°
            elif angle_diff < 7 * np.pi / 8:
                action[1] = 7  # 67.5°
            else:
                action[1] = 8  # 90°（最大右转角度）
        else:
            # 需要向左转向
            if angle_diff > -np.pi / 8:
                action[1] = 4  # 0°（不需要转向）
            elif angle_diff > -3 * np.pi / 8:
                action[1] = 3  # -22.5°
            elif angle_diff > -5 * np.pi / 8:
                action[1] = 2  # -45°
            elif angle_diff > -7 * np.pi / 8:
                action[1] = 1  # -67.5°
            else:
                action[1] = 0  # -90°（最大左转角度）
        
        # 始终选择最大加速度，全力逃脱
        action[0] = 4
    
    return np.array(action, dtype=np.int64)

def rule_based_decision(obs, agent_type):
    if '猎手' in agent_type:
        return hunter_decision(obs)
    elif '逃脱者' in agent_type:
        return escaper_decision(obs)
    else:
        raise ValueError("Unknown agent type")